/*
 * Copyright (c) 2021 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef ECMASCRIPT_COMPILER_LLVM_CODEGEN_H
#define ECMASCRIPT_COMPILER_LLVM_CODEGEN_H

#include <iostream>
#include <list>
#include <map>
#include <sys/mman.h>
#include <vector>

#include "ecmascript/compiler/binary_section.h"
#include "ecmascript/compiler/code_generator.h"
#include "ecmascript/compiler/compiler_log.h"
#include "ecmascript/compiler/llvm_ir_builder.h"
#include "ecmascript/ecma_macros.h"
#include "ecmascript/js_thread.h"

#if defined(__clang__)
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wshadow"
#pragma clang diagnostic ignored "-Wunused-parameter"
#elif defined(__GNUC__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wshadow"
#pragma GCC diagnostic ignored "-Wunused-parameter"
#endif

#include "ecmascript/mem/machine_code.h"
#include "ecmascript/mem/region.h"
#include "libpandabase/utils/asan_interface.h"
#include "llvm-c/Analysis.h"
#include "llvm-c/Core.h"
#include "llvm-c/ExecutionEngine.h"
#include "llvm-c/Target.h"
#include "llvm-c/Transforms/PassManagerBuilder.h"
#include "llvm-c/Transforms/Scalar.h"
#include "llvm-c/Types.h"
#include "llvm/ExecutionEngine/Interpreter.h"
#include "llvm/IR/Instructions.h"
#include "llvm/ExecutionEngine/SectionMemoryManager.h"
#include "llvm/ExecutionEngine/MCJIT.h"
#include "llvm/Support/Host.h"

#if defined(__clang__)
#pragma clang diagnostic pop
#elif defined(__GNUC__)
#pragma GCC diagnostic pop
#endif

namespace panda::ecmascript::kungfu {
struct CodeInfo {
    using sectionInfo = std::pair<uint8_t *, size_t>;
    CodeInfo()
    {
        continuousSectionsMem_ = static_cast<uint8_t *>(mmap(nullptr, MAX_MACHINE_CODE_SIZE, protRWX, flags, -1, 0));
        if (continuousSectionsMem_ == reinterpret_cast<uint8_t *>(-1)) {
            continuousSectionsMem_ = nullptr;
        }
        if (continuousSectionsMem_ != nullptr) {
            ASAN_UNPOISON_MEMORY_REGION(continuousSectionsMem_, MAX_MACHINE_CODE_SIZE);
        }
        // align machineCode for aarch64
        continuousSectionsMem_ += MachineCode::DATA_OFFSET +
            AlignUp(sizeof(Region), static_cast<size_t>(MemAlignment::MEM_ALIGN_REGION));
        secInfos_.fill(std::make_pair(nullptr, 0));
    }
    ~CodeInfo()
    {
        Reset();
        if (continuousSectionsMem_ != nullptr) {
            ASAN_POISON_MEMORY_REGION(continuousSectionsMem_, MAX_MACHINE_CODE_SIZE);
            munmap(continuousSectionsMem_, MAX_MACHINE_CODE_SIZE);
        }
        continuousSectionsMem_ = nullptr;
    }

    uint8_t *Alloca(uintptr_t size)
    {
        // align up for rodata section
        size = AlignUp(size, static_cast<size_t>(MemAlignment::MEM_ALIGN_REGION));
        uint8_t *addr = nullptr;
        if (bufferPos_ + size > MAX_MACHINE_CODE_SIZE) {
            LOG_COMPILER(ERROR) << std::hex << "Alloca Section failed. Current bufferPos_:" << bufferPos_
                      << " plus size:" << size << "exceed MAX_MACHINE_CODE_SIZE:" << MAX_MACHINE_CODE_SIZE;
            return nullptr;
        }
        addr = continuousSectionsMem_ + bufferPos_;
        bufferPos_ += size;
        return addr;
    }

    uint8_t *AllocaCodeSection(uintptr_t size, const char *sectionName)
    {
        uint8_t *addr = Alloca(size);
        auto curSec = ElfSection(sectionName);
        codeInfo_.push_back({addr, size});
        if (curSec.isValidAOTSec()) {
            secInfos_[curSec.GetIntIndex()] = std::make_pair(addr, size);
        }
        return addr;
    }

    uint8_t *AllocaDataSection(uintptr_t size, const char *sectionName)
    {
        uint8_t *addr = nullptr;
        auto curSec = ElfSection(sectionName);
        if (curSec.isSequentialAOTSec()) {
            addr = Alloca(size);
        } else { // the assumption here is llvm backend wouldn't emit a section with duplicated name
            discreteSectionsAllocator_.insert(make_pair(sectionName, std::vector<uint8_t>()));
            ASSERT(discreteSectionsAllocator_.find(sectionName) != discreteSectionsAllocator_.end());
            discreteSectionsAllocator_[sectionName].resize(size);
            addr = static_cast<uint8_t *>(discreteSectionsAllocator_[sectionName].data());
        }
        if (curSec.isValidAOTSec()) {
            secInfos_[curSec.GetIntIndex()] = std::make_pair(addr, size);
        }
        return addr;
    }

    void Reset()
    {
        codeInfo_.clear();
        discreteSectionsAllocator_.clear();
        continuousSectionsMem_ = nullptr;
        bufferPos_ = 0;
    }

    uint8_t *GetSectionAddr(ElfSecName sec) const
    {
        auto curSection = ElfSection(sec);
        auto idx = curSection.GetIntIndex();
        return const_cast<uint8_t *>(secInfos_[idx].first);
    }

    size_t GetSectionSize(ElfSecName sec) const
    {
        auto curSection = ElfSection(sec);
        auto idx = curSection.GetIntIndex();
        return secInfos_[idx].second;
    }

    std::vector<std::pair<uint8_t *, uintptr_t>> GetCodeInfo() const
    {
        return codeInfo_;
    }

    template <class Callback>
    void IterateSecInfos(const Callback &cb) const
    {
        for (size_t i = 0; i < secInfos_.size(); i++) {
            if (secInfos_[i].second == 0) {
                continue;
            }
            cb(i, secInfos_[i]);
        }
    }

private:
    static constexpr size_t MAX_MACHINE_CODE_SIZE = (1 << 20);  // 1M
    static constexpr int protRWX = PROT_READ | PROT_WRITE | PROT_EXEC;  // NOLINT(hicpp-signed-bitwise)
    static constexpr int flags = MAP_ANONYMOUS | MAP_SHARED;            // NOLINT(hicpp-signed-bitwise)
    uint8_t *continuousSectionsMem_ {nullptr};
    size_t bufferPos_ {0};
    std::array<sectionInfo, static_cast<int>(ElfSecName::SIZE)> secInfos_;
    std::map<std::string, std::vector<uint8_t>> discreteSectionsAllocator_;
    std::vector<std::pair<uint8_t *, uintptr_t>> codeInfo_ {}; // info for disasssembler, planed to be deprecated
};

struct LOptions {
    uint32_t optLevel : 2; // 2 bits for optimized level 0-4
    uint32_t genFp : 1; // 1 bit for whether to generated frame pointer or not
    uint32_t relocMode : 3; // 3 bits for relocation mode
    // 3: default optLevel, 1: generating fp, 2: PIC mode
    LOptions() : optLevel(3), genFp(1), relocMode(2) {};
    LOptions(size_t level, bool genFp, size_t relocMode) : optLevel(level), genFp(genFp), relocMode(relocMode) {};
};

class LLVMAssembler {
public:
    explicit LLVMAssembler(LLVMModuleRef module, LOptions option = LOptions());
    virtual ~LLVMAssembler();
    void Run();
    const LLVMExecutionEngineRef &GetEngine()
    {
        return engine_;
    }
    void Disassemble(const std::map<uintptr_t, std::string> &addr2name, const CompilerLog &log) const;
    static int GetFpDeltaPrevFramSp(LLVMValueRef fn, const CompilerLog &log);
    static void Disassemble(uint8_t *buf, size_t size);

    uintptr_t GetSectionAddr(ElfSecName sec) const
    {
        return reinterpret_cast<uintptr_t>(codeInfo_.GetSectionAddr(sec));
    }

    uint32_t GetSectionSize(ElfSecName sec) const
    {
        return static_cast<uint32_t>(codeInfo_.GetSectionSize(sec));
    }

    void *GetFuncPtrFromCompiledModule(LLVMValueRef function)
    {
        return LLVMGetPointerToGlobal(engine_, function);
    }

    template <class Callback>
    void IterateSecInfos(const Callback &cb) const
    {
        codeInfo_.IterateSecInfos(cb);
    }

private:
    void UseRoundTripSectionMemoryManager();
    bool BuildMCJITEngine();
    void BuildAndRunPasses();
    void Initialize(LOptions option);

    LLVMMCJITCompilerOptions options_ {};
    LLVMModuleRef module_;
    LLVMExecutionEngineRef engine_ {nullptr};
    char *error_ {nullptr};
    struct CodeInfo codeInfo_ {};
};

class LLVMIRGeneratorImpl : public CodeGeneratorImpl {
public:
    explicit LLVMIRGeneratorImpl(LLVMModule *module, bool enableLog)
        : module_(module), enableLog_(enableLog) {}
    ~LLVMIRGeneratorImpl() = default;
    void GenerateCodeForStub(Circuit *circuit, const ControlFlowGraph &graph, size_t index,
                             const CompilationConfig *cfg) override;
    void GenerateCode(Circuit *circuit, const ControlFlowGraph &graph, const CompilationConfig *cfg,
        const JSMethod *method) override;

    bool IsLogEnabled() const
    {
        return enableLog_;
    }

private:
    LLVMModule *module_;
    bool enableLog_ {false};
};
}  // namespace panda::ecmascript::kungfu
#endif  // ECMASCRIPT_COMPILER_LLVM_CODEGEN_H
